package com.bytes32.stats;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.nio.DoubleBuffer;
import java.util.Collection;
import java.util.Collections;
import java.util.Deque;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Random;

import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableUtils;

/**
 * Creates a structure which
 * is able to estimate with limited memory
 * the q-th order statistics for large data
 * sets and with unknown number of elements
 * 
 * This class is not thread safe
 * if you must use multiple threads
 * use appropriate locking
 * 
 * @author murariuf
 *
 */
public class LargeQuantiles implements Writable {

	/*
	 * Number of buffers
	 */
	private int b;
	
	/*
	 * Sampling rate also the starting weight of the new buffers
	 */
	private int r = 1;
	
	/*
	 * Size of the buffers
	 */
	private int k;
	
	/*
	 * height of the tree where sampling starts
	 */
	private int h; 
	
	/*
	 * Current height of the tree
	 */
	private int H = 0;
	
	/*
	 * Starting level for new buffers
	 */
	private int l = 0;
	
	/*
	 * The list of buffers
	 */
	protected Deque<QBuffer> buffers;

	/* used for random sampling */
	private Random rand = new Random();
	
	/*
	 * Where stuff gets added
	 */
	protected QBuffer current;
	
	/**
	 * Create a q-th order statistics estimator
	 * for large data sets
	 * 
	 * @param startSampling
	 *  level at which you can start sampling
	 *  minimum 3
	 * @param conf
	 *  {@link Config._01} lowest precision
	 *  {@link Config._15} highest precision
	 */
	public LargeQuantiles(final int startSampling, Config conf) {
		if (startSampling < 2)
			throw new IllegalArgumentException("Minimum number of buffers must be 2");
		if (conf == null){
			throw new IllegalArgumentException("Must specify a configuration");
		}
		h = startSampling;
		b = conf.getB();
		k = conf.getK();
		/*Init the buffers*/
		buffers = new LinkedList<QBuffer>();
		/* create a new buffer in the list*/
		newCurrentBuffer();
	}
	
	/**
	 * Adds a buffer to be processed in r increments
	 * no inner buffer support, bring your own
	 * 
	 * @param buffer
	 * @throws IllegalArgumentException
	 * 	when input is null or empty
	 */
	public void add(double...input){
		this._add(input);
	}

	/**
	 * Adds a double[] to the buffer with sampling 
	 * @param input
	 */
	protected void _add(double... input){
		if (input == null || input.length == 0)
			throw new IllegalArgumentException("No quantiles to estimate");
		double value = 0;
		DoubleBuffer buffer = DoubleBuffer.wrap(input);
		while (buffer.hasRemaining()){
			if (r > 1){
				int chunk_size = buffer.position()+r < input.length?r:input.length-buffer.position();
				double[] chunk = new double[chunk_size];
				buffer.get(chunk);
				value = chunk[rand.nextInt(chunk_size)];
			} else {
				value = buffer.get();
			}
			addInternal(value);
		}
	}
	
	/**
	 * This is where magic happens
	 * @param value
	 */
	private void addInternal(double value){
		BufferStatus stat = current.add(value);		
		if (stat == BufferStatus.FULL && buffers.size() < b){
			/* sort the buffer once we are done with it
			 * no need to sort buffers in collapse */
			current.sort(); 
			newCurrentBuffer();
			return;
		} else if (stat == BufferStatus.PARTIAL){
			/* nothing to see, go away */
			return;
		}
		/* time for collapse, sort the current buffer*/
		current.sort(); 

		/*We now have all full buffers and the buffer count has reached b*/
		int minLevel = current.level();
		int initialSize = buffers.size();
		
		Collection<QBuffer> collapseBufers = new LinkedList<QBuffer>();
		
		Iterator<? extends QBuffer> descIter = buffers.descendingIterator();
		while(descIter.hasNext()){
			/* pick the buffers to collapse */
			QBuffer buf = descIter.next();
			if (buf.level() == minLevel || (buf.level() > minLevel && collapseBufers.size() == 1)){
				/* If we only have one buffer at the minimum level increase the minimum level */
				if (buf.level() > minLevel && collapseBufers.size() == 1){
					minLevel = buf.level();
					current.level(minLevel);
				}
				/*Remove from the list of buffers & add to the collection to be compacted*/
				descIter.remove();
				collapseBufers.add(buf);
			} 
		}
		QBuffer result = MergeAndSelect.collapse(collapseBufers, ++minLevel);
		buffers.add(result);
		/* Once sampling has started H >= h
		 * If we collapsed all the buffers then we double r 
		 * thus decreasing the sampling rate and the height of the
		 * tree and the level of the new buffers 
		 */ 
		if (collapseBufers.size() == initialSize){
			if ( H >= h){
				r*=2;
				l++; /* buffers only increase in level once sampling starts */
			}
			H=result.level(); /* new height of the tree */
		}
		/* Create and empty current buffer */
		newCurrentBuffer();
		/*All happy now */
		addInternal(value);
	}
	
	/**
	 * Returns the desired quantiles, however this
	 * is faster but more un-balanced towards extremities
	 * 
	 * @param q
	 * @return
	 */
	public double[] naiveQuantile(double... q){
		return MergeAndSelect.naiveOutput(buffers, q);
	}
	
	/**
	 * Returns the desired quantiles, however this is slower
	 * but more precise.
	 * @param q
	 * @return
	 */
	public double[] quantile(double... q){
		return MergeAndSelect.output(buffers, q);
	}
	
	/* compacts the buffers into maximum 2 buffers, one full and one partial */
	public LargeQuantiles compact(){
		this.buffers = new LinkedList<QBuffer>(MergeAndSelect.preReduce(buffers)); 
		return this;
	}
	
	/**
	 * @return
	 */
	public Collection<QBuffer> buffers(){
		return Collections.unmodifiableCollection(this.buffers);
	}
	
	private void newCurrentBuffer(){
		current = new InputBuffer(new double[k], r, l, BufferStatus.EMPTY);
		buffers.add(current);
	}

	/**
	 * Gets the number of buffers
	 * @return
	 */
	public int getB() {
		return b;
	}

	/**
	 * Gets the height of the tree before sampling starts
	 * @return
	 */
	public int getH() {
		return h;
	}

	/**
	 * Gets the sampling rate
	 * @return
	 */
	public int getR() {
		return r;
	}

	/**
	 * Sets the sampling rate
	 * @param r
	 * @throws IllegalArgumentException if r < 1
	 */
	protected void setR(int r) {
		if (r < 1)
			throw new IllegalArgumentException("Sampling rate cannot be less than 1");
		this.r = r;
	}

	@Override
	public void readFields(DataInput stream) throws IOException {
		this.h = WritableUtils.readVInt(stream);
		this.b = WritableUtils.readVInt(stream);
		this.k = WritableUtils.readVInt(stream);
		this.H = WritableUtils.readVInt(stream);
		this.l = WritableUtils.readVInt(stream);
		this.r = WritableUtils.readVInt(stream);
		int size = WritableUtils.readVInt(stream);
		this.buffers = new LinkedList<QBuffer>();
		for (int i = 0; i < size; i++){
			QBuffer buf = new InputBuffer(0, 0, 0, BufferStatus.EMPTY);
			buf.readFields(stream);
			this.buffers.add(buf);
		}
	}

	@Override
	public void write(DataOutput stream) throws IOException {
		WritableUtils.writeVInt(stream, this.h);
		WritableUtils.writeVInt(stream, this.b);
		WritableUtils.writeVInt(stream, this.k);
		WritableUtils.writeVInt(stream, this.H);
		WritableUtils.writeVInt(stream, this.l);
		WritableUtils.writeVInt(stream, this.r);
		WritableUtils.writeVInt(stream, this.buffers.size());
		Iterator<QBuffer> iter = this.buffers.iterator();
		while (iter.hasNext()){
			QBuffer buf = iter.next();
			buf.write(stream);
		}
	}

	public int getK() {
		return k;
	}
		
}
